from unsloth import FastLanguageModel
from datasets import load_dataset, concatenate_datasets
import torch
from transformers import TrainingArguments
from trl import SFTTrainer
from unsloth import is_bfloat16_supported

# 加载两个数据集（如您截图所示）
reasoning_dataset = load_dataset("unsloth/OpenMathReasoning-mini", split="cot")
non_reasoning_dataset = load_dataset("mlabonne/FineTome-100k", split="train")

print(f"推理数据集大小: {len(reasoning_dataset)}")
print(f"非推理数据集大小: {len(non_reasoning_dataset)}")

# 采样平衡数据集（避免非推理数据主导）
max_reasoning_samples = min(len(reasoning_dataset), 50000)
max_non_reasoning_samples = min(len(non_reasoning_dataset), 50000)

reasoning_subset = reasoning_dataset.select(range(max_reasoning_samples))
non_reasoning_subset = non_reasoning_dataset.select(range(max_non_reasoning_samples))

# 合并数据集
combined_dataset = concatenate_datasets([reasoning_subset, non_reasoning_subset])
combined_dataset = combined_dataset.train_test_split(test_size=0.1, seed=42)

print(f"合并后训练集大小: {len(combined_dataset['train'])}")
print(f"合并后测试集大小: {len(combined_dataset['test'])}")